(EAI-990): Refactor search as a tool #705

mongodben · 2025-05-09T20:46:40Z

Jira: https://jira.mongodb.org/browse/EAI-990

Changes

Refactor chatbot to use tool calling for search
- general generateResponseWithSearchTool() to be used with arbitrary search tool in mongodb-chatbot-server
Add our implementation of search as tool to chatbot-server-mongodb-public
- as a result of this change, can remove 2 LLM preprocessors, the step back query generator, and also the metadata extractor
Upgrade to use GPT-4.1 in chatbot
Fix broken tests in branch

Notes

packages/mongodb-chatbot-server/src/processors/generateResponseWithSearchTool.ts

nlarew

LGTM - lots going on in this one so generally trusting you that it works as expected. Proof is in the pudding!

packages/chatbot-server-mongodb-public/src/config.ts

nlarew · 2025-06-03T21:08:08Z

packages/chatbot-server-mongodb-public/src/processors/makeMongoDbReferences.test.ts

+    const result = makeMongoDbReferences(
+      chunks.map((c) => ({
+        ...c,
+        title: c.metadata?.pageTitle,
+      }))
+    );


Why do we need to modify the chunks here?

umm we probably dont. esp based on what the change is doing. will poke at and either 1) report back on why we do, or 2) revert

update: reverting. no change in behavior

nlarew · 2025-06-03T21:14:28Z

packages/chatbot-server-mongodb-public/src/systemPrompt.ts

+  'If the user query is already a "good" search query, do not modify it.',
+  'For one word queries like "or", "and", "exists", if the query corresponds to a MongoDB operation, transform it into a fully formed question. Ex: If the user query is "or", transform it into "what is the $or operator in MongoDB?".',
+  "You should also transform the user query into a fully formed question, if relevant.",
+  `Only generate ONE ${SEARCH_TOOL_NAME} tool call per user message unless there are clearly multiple distinct queries needed to answer the user query.`,


I assume we handle the multiple tool call case in regards to reference links?

yea, it all appends to an array of references, which is streamed after the text generation

nlarew · 2025-06-03T21:17:41Z

packages/chatbot-server-mongodb-public/src/tracing/extractTracingData.ts

@@ -12,6 +12,7 @@ export function extractTracingData(
  messages: Message[],
  assistantMessageId: ObjectId
 ) {
+  // FIXME: this is throwing after the generation is complete. don't forget to fix before merge of EAI-990


Is this and other TODO in this file fixed in another PR?

fixed in other PR! #759

nlarew · 2025-06-03T21:20:54Z

packages/mongodb-chatbot-server/src/processors/generateResponseWithSearchTool.test.ts

+          dataStreamer,
+        });
+
+        // TODO: verify dataStreamer was called


Need todo still or covered in another PR?

hmm not sure, will add here just in case. easy fix.

was added in later PR but added here as well, just in case

packages/mongodb-chatbot-server/src/processors/makeDefaultReferenceLinks.ts

Co-authored-by: Nick Larew <nick.larew@mongodb.com>

* (EAI-988): Refactor `GenerateResponse` for tool call support (#687) * refactor GenerateRespose * Clean up imports * consolidate generate user prompt to the legacy file * update test config imports * Fix broken tests * (EAI-989): Refactor verified answers to wrap `GenerateResponse` (#688) verified answer generate response Co-authored-by: Ben Perlmutter <mongodben@mongodb.com> * handle streaming * separate generateresponse * typo fix --------- Co-authored-by: Ben Perlmutter <mongodben@mongodb.com> * (EAI-990): Refactor search as a tool (#705) * refactor GenerateRespose * Clean up imports * consolidate generate user prompt to the legacy file * update test config imports * Fix broken tests * get started * nominally working generate res w/ search * small refactors * aint pretty but fully functional * hacky if more functional * more hack * tools * functional if not pretty * Add processing * working tool calling * making progress * keepin on * Clean config * working e2e * update model version * Remove no longer used stuff * decouple search results for references and whats shown to model * fix scripts build errs * fix broken tests * update default ref links * fix broken tests * Apply suggestions from code review Co-authored-by: Nick Larew <nick.larew@mongodb.com> * revert default reference links * adding missing test --------- Co-authored-by: Ben Perlmutter <mongodben@mongodb.com> Co-authored-by: Nick Larew <nick.larew@mongodb.com> * (EAI-992): Remove `ChatLlm` (#751) * refactor GenerateRespose * Clean up imports * consolidate generate user prompt to the legacy file * update test config imports * Fix broken tests * get started * nominally working generate res w/ search * small refactors * aint pretty but fully functional * hacky if more functional * more hack * tools * functional if not pretty * Add processing * working tool calling * making progress * keepin on * Clean config * working e2e * update model version * Remove no longer used stuff * decouple search results for references and whats shown to model * fix scripts build errs * remove ChatLlm * lite fixes * Remove stub --------- Co-authored-by: Ben Perlmutter <mongodben@mongodb.com> * (EAI-993): deprecate framework (#752) * refactor GenerateRespose * Clean up imports * consolidate generate user prompt to the legacy file * update test config imports * Fix broken tests * get started * nominally working generate res w/ search * small refactors * aint pretty but fully functional * hacky if more functional * more hack * tools * functional if not pretty * Add processing * working tool calling * making progress * keepin on * Clean config * working e2e * update model version * Remove no longer used stuff * decouple search results for references and whats shown to model * fix scripts build errs * fix broken tests * deprecation * build out docs following last week convo * clean up spec + contact * fix merge funk * docs updates --------- Co-authored-by: Ben Perlmutter <mongodben@mongodb.com> * (EAI-1071): Fix broken Atlas OpenAPI ingest (#765) * update to fix broken test * Update packages/ingest-mongodb-public/src/sources/snooty/snootyAstToOpenApiSpec.ts --------- Co-authored-by: Ben Perlmutter <mongodben@mongodb.com> * (EAI-995): add guardrail (#755) * refactor GenerateRespose * Clean up imports * consolidate generate user prompt to the legacy file * update test config imports * Fix broken tests * get started * nominally working generate res w/ search * small refactors * aint pretty but fully functional * hacky if more functional * more hack * tools * functional if not pretty * Add processing * working tool calling * making progress * keepin on * Clean config * working e2e * update model version * Remove no longer used stuff * decouple search results for references and whats shown to model * fix scripts build errs * fix broken tests * update default ref links * fix broken tests * input guardrail refactor * guardrail works well * simpler validity metric * add guardrail to server * add next step todo * llm refusal msg * remove TODO comment * merge fix * fix unnec changes * NL feedback --------- Co-authored-by: Ben Perlmutter <mongodben@mongodb.com> * fix type in text * (EAI-991 & EAI-1050): Evaluate and clean up retrieval as a tool (#757) * refactor GenerateRespose * Clean up imports * consolidate generate user prompt to the legacy file * update test config imports * Fix broken tests * get started * nominally working generate res w/ search * small refactors * aint pretty but fully functional * hacky if more functional * more hack * tools * functional if not pretty * Add processing * working tool calling * making progress * keepin on * Clean config * working e2e * update model version * Remove no longer used stuff * decouple search results for references and whats shown to model * fix scripts build errs * fix broken tests * update default ref links * fix broken tests * input guardrail refactor * guardrail works well * simpler validity metric * add guardrail to server * add next step todo * llm refusal msg * remove TODO comment * evals on new architecture * Get urls in a way that supports verified answers * dont eval on retrieved elems if no context * Cleaner handling * update trace handling * update trace handling * undo git funk * handle undefined case * Fix tracing test --------- Co-authored-by: Ben Perlmutter <mongodben@mongodb.com> * remove console logs + redunancies --------- Co-authored-by: Ben Perlmutter <mongodben@mongodb.com> Co-authored-by: Nick Larew <nick.larew@mongodb.com>

Ben Perlmutter added 13 commits April 28, 2025 12:08

refactor GenerateRespose

2e286c3

Clean up imports

0be9fe2

consolidate generate user prompt to the legacy file

2556bc2

update test config imports

feee66a

Fix broken tests

f6fe862

get started

0fc58bd

nominally working generate res w/ search

ea40c6d

small refactors

c3e69e3

aint pretty but fully functional

0345453

hacky if more functional

a4144db

more hack

04076d2

tools

8372bd3

functional if not pretty

24d1cf7

mongodben changed the title ~~Eai 990: functional tool calling~~ Eai 990: tool calling search May 9, 2025

Ben Perlmutter added 11 commits May 12, 2025 09:20

Add processing

40fa9d1

working tool calling

31bd0a8

making progress

653aa59

keepin on

9d68d50

Clean config

57e65a2

working e2e

aefa9ca

update model version

c33f64e

Merge remote-tracking branch 'upstream/retrieval_tool_call' into EAI-990

c446756

Remove no longer used stuff

728416f

decouple search results for references and whats shown to model

479ccb8

fix scripts build errs

cc8dd45

mongodben mentioned this pull request May 27, 2025

(EAI-992): Remove ChatLlm #751

Merged

mongodben changed the title ~~Eai 990: tool calling search~~ (EAI-990): Refactor search as a tool May 27, 2025

Ben Perlmutter added 3 commits May 27, 2025 15:46

fix broken tests

e768dd6

update default ref links

2bce005

fix broken tests

b8f3754

mongodben marked this pull request as ready for review May 27, 2025 20:52

mongodben mentioned this pull request May 29, 2025

(EAI-995): add guardrail #755

Merged

nlarew reviewed May 30, 2025

View reviewed changes

packages/mongodb-chatbot-server/src/processors/generateResponseWithSearchTool.ts Show resolved Hide resolved

nlarew approved these changes Jun 3, 2025

View reviewed changes

mongodben and others added 3 commits June 4, 2025 08:48

Apply suggestions from code review

298d01d

Co-authored-by: Nick Larew <nick.larew@mongodb.com>

revert default reference links

40c9204

adding missing test

bb7d0e1

mongodben merged commit 50d3f44 into retrieval_tool_call Jun 4, 2025
1 check was pending

mongodben deleted the EAI-990 branch June 4, 2025 13:13

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

(EAI-990): Refactor search as a tool #705

(EAI-990): Refactor search as a tool #705

Uh oh!

mongodben commented May 9, 2025 •

edited

Loading

Uh oh!

Uh oh!

nlarew left a comment

Uh oh!

Uh oh!

nlarew Jun 3, 2025

Uh oh!

mongodben Jun 4, 2025

Uh oh!

mongodben Jun 4, 2025

Uh oh!

nlarew Jun 3, 2025

Uh oh!

mongodben Jun 4, 2025

Uh oh!

nlarew Jun 3, 2025

Uh oh!

mongodben Jun 4, 2025

Uh oh!

nlarew Jun 3, 2025

Uh oh!

mongodben Jun 4, 2025

Uh oh!

mongodben Jun 4, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

(EAI-990): Refactor search as a tool #705

(EAI-990): Refactor search as a tool #705

Uh oh!

Conversation

mongodben commented May 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Changes

Notes

Uh oh!

Uh oh!

nlarew left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mongodben commented May 9, 2025 •

edited

Loading